use std::fs::File;
use std::io::{BufReader, Read};
use std::rc::Rc;
use crate::Tensor;

pub fn load_digits() -> (Tensor<f64>, Tensor<f64>, Tensor<f64>, Tensor<f64>) {
    let file = File::open("/Users/khlevnov/twinkle/python/datasets/digits.bin").unwrap();
    let mut buf_reader = BufReader::new(file);

    let mut floats = vec![];
    let bytes_count = 1797 * (64 + 10);

    let mut buf = [0u8; 8];
    for _ in 0..bytes_count {
        buf_reader.read_exact(&mut buf).unwrap();
        floats.push(f64::from_le_bytes(buf));
    }

    let x_train_data = floats.iter()
        .take(1000 * 64)
        .copied()
        .collect::<Box<[f64]>>();
    let x_train: Tensor<f64> = (Rc::new(x_train_data), [1000, 64].into()).into();

    let x_test_data = floats.iter()
        .skip(1000 * 64)
        .take(797 * 64)
        .copied()
        .collect::<Box<[f64]>>();
    let x_test: Tensor<f64> = (Rc::new(x_test_data), [797, 64].into()).into();

    let y_train_data = floats.iter()
        .skip(1797 * 64)
        .take(1000 * 10)
        .copied()
        .collect::<Box<[f64]>>();
    let y_train: Tensor<f64> = (Rc::new(y_train_data), [1000, 10].into()).into();

    let y_test_data = floats.iter()
        .skip(1797 * 64 + 1000 * 10)
        .take(797 * 10)
        .copied()
        .collect::<Box<[f64]>>();
    let y_test: Tensor<f64> = (Rc::new(y_test_data), [797, 10].into()).into();

    (x_train, x_test, y_train, y_test)
}
